use crate::Var;
use crate::nn::active::Activation;
use crate::nn::Neural;

#[derive(Debug)]
pub struct Linear {
    neurons: Vec<Neural>,
}

impl Linear {
    /// `output_size` neurons, each neuron has `input_size`
    pub fn new(input_size: usize, output_size: usize, activation: Option<Activation>) -> Self {
        Self {
            neurons: (0..output_size).into_iter().map(|_| Neural::new(input_size, activation)).collect()
        }
    }

    pub fn forward(&self, xs: &[Var]) -> Vec<Var> {
        self.neurons.iter().map(|neuron| neuron.forward(xs)).collect()
    }

    pub fn parameters(&self) -> impl Iterator<Item = &Var> {
        self.neurons.iter().flat_map(|neural| neural.parameters())
    }
}